from __future__ import annotations

from dataclasses import dataclass, field
from typing import Iterable

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity


@dataclass
class MemoryDocument:
    content: str
    tags: dict[str, str] = field(default_factory=dict)


class MemoryStore:
    """Simple TF-IDF backed memory store for prototype purposes."""

    def __init__(self) -> None:
        self._documents: list[MemoryDocument] = []
        self._vectorizer = TfidfVectorizer(stop_words="english")
        self._matrix: np.ndarray | None = None

    def add(self, document: MemoryDocument) -> None:
        self._documents.append(document)
        self._rebuild_index()

    def bulk_add(self, documents: Iterable[MemoryDocument]) -> None:
        self._documents.extend(documents)
        self._rebuild_index()

    def search(self, query: str, top_k: int = 3) -> list[MemoryDocument]:
        if not self._documents:
            return []
        if self._matrix is None:
            self._rebuild_index()
            if self._matrix is None:
                return []
        query_vec = self._vectorizer.transform([query])
        scores = cosine_similarity(query_vec, self._matrix).flatten()
        top_indices = scores.argsort()[::-1][:top_k]
        return [self._documents[i] for i in top_indices if scores[i] > 0]

    def _rebuild_index(self) -> None:
        if not self._documents:
            self._matrix = None
            return
        # Filter out empty documents and ensure non-empty corpus
        corpus = [doc.content if doc.content.strip() else "empty_document"
                  for doc in self._documents]
        # Additional safety check to prevent empty vocabulary error
        if all(not text.strip() for text in corpus):
            corpus = ["placeholder_document"]  # Ensure at least one non-empty document
        self._matrix = self._vectorizer.fit_transform(corpus)


__all__ = ["MemoryStore", "MemoryDocument"]
